import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import math
import copy
from torch.utils.data import Dataset, DataLoader

# -------------------------- 1. 基础配置（设备、随机种子） --------------------------
seed = 42
torch.manual_seed(seed)
if torch.cuda.is_available():
    torch.cuda.manual_seed_all(seed)

Device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using Device: {Device}")


# -------------------------- 2. 用户指定的SimpleDataset --------------------------
class SimpleDataset(Dataset):
    def __init__(self, num):
        prefix = torch.tensor([0])  # 起始符号0
        suffix = torch.tensor([9])  # 结束符号9
        self.data = [torch.randint(1, 9, size=(5,), dtype=torch.long) for _ in range(num)]
        self.labels = [torch.cat([prefix, _, suffix], dim=0) for _ in copy.deepcopy(self.data)]

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx], self.labels[idx]  # data: (5,), label: (7,)


def create_data(batch_size=5, num=100):
    dataset = SimpleDataset(num)
    data_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
    return data_loader


# -------------------------- 3. Transformer核心模块（自定义编码器替换关键部分） --------------------------
# 3.1 位置编码（不变）
class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=10, dropout=0.1):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)

        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))

        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(1)  # (max_len, 1, d_model)
        self.register_buffer('pe', pe)

    def forward(self, x):
        # x: (batch_size, seq_len, d_model)
        x = x + self.pe[:x.size(0), :, :]
        return self.dropout(x)


# 3.2 自定义EncoderLayer（不变，已用自定义逻辑）
class EncoderLayer(nn.Module):
    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1):
        super(EncoderLayer, self).__init__()
        # 多头注意力仍用PyTorch实现（若需完全自定义可替换，但此处先保证收敛，重点替换编码器容器）
        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim_feedforward, d_model)

        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
        self.activation = F.relu

    def forward(self, src, src_mask=None, src_key_padding_mask=None):
        # src: (seq_len, batch_size, d_model)
        src2 = self.self_attn(
            src, src, src,
            attn_mask=src_mask,
            key_padding_mask=src_key_padding_mask
        )[0]  # 注意力输出（取元组第0个元素）
        src = src + self.dropout1(src2)  # 残差连接
        src = self.norm1(src)  # 层归一化

        # 前馈网络
        src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))
        src = src + self.dropout2(src2)  # 残差连接
        src = self.norm2(src)  # 层归一化
        return src


# -------------------------- 关键：自定义TransformerEncoder（替换nn.TransformerEncoder） --------------------------
class CustomTransformerEncoder(nn.Module):
    def __init__(self, encoder_layer, num_layers, norm=None):
        """
        与PyTorch自带nn.TransformerEncoder接口完全对齐：
        - encoder_layer: 单个编码器层（如自定义的EncoderLayer）
        - num_layers: 编码器层数
        - norm: 输出层归一化（可选，默认None）
        """
        super(CustomTransformerEncoder, self).__init__()
        # 用ModuleList堆叠num_layers个编码器层（确保参数可训练）
        self.layers = nn.ModuleList([copy.deepcopy(encoder_layer) for _ in range(num_layers)])
        self.num_layers = num_layers
        self.norm = norm  # 可选的输出归一化

    def forward(self, src, mask=None, src_key_padding_mask=None):
        """
        与PyTorch自带接口对齐，前向传播逻辑一致：
        1. 依次通过每一层编码器
        2. 若有norm，最后对输出做归一化
        - src: (seq_len, batch_size, d_model) → 输入序列
        - mask: (seq_len, seq_len) → 自注意力掩码（可选）
        - src_key_padding_mask: (batch_size, seq_len) → 填充掩码（可选）
        - 返回: (seq_len, batch_size, d_model) → 编码器最终输出（memory）
        """
        output = src  # 初始输入为src

        # 循环通过每一层编码器
        for layer in self.layers:
            output = layer(
                output,
                src_mask=mask,
                src_key_padding_mask=src_key_padding_mask
            )

        # 若有输出归一化，应用归一化
        if self.norm is not None:
            output = self.norm(output)

        return output


# 3.3 自定义DecoderLayer（不变，兼容自定义编码器输出）
class DecoderLayer(nn.Module):
    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1):
        super(DecoderLayer, self).__init__()
        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
        self.multihead_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)

        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim_feedforward, d_model)

        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
        self.dropout3 = nn.Dropout(dropout)
        self.activation = F.relu

    def forward(self, tgt, memory, tgt_mask=None, memory_mask=None,
                tgt_key_padding_mask=None, memory_key_padding_mask=None):
        # tgt: (tgt_seq_len, batch_size, d_model)；memory: 自定义编码器输出（同格式）
        tgt2 = self.self_attn(tgt, tgt, tgt, attn_mask=tgt_mask, key_padding_mask=tgt_key_padding_mask)[0]
        tgt = tgt + self.dropout1(tgt2)
        tgt = self.norm1(tgt)

        # 交叉注意力：使用自定义编码器的输出memory
        tgt2 = self.multihead_attn(tgt, memory, memory, attn_mask=memory_mask, key_padding_mask=memory_key_padding_mask)[0]
        tgt = tgt + self.dropout2(tgt2)
        tgt = self.norm2(tgt)

        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
        tgt = tgt + self.dropout3(tgt2)
        tgt = self.norm3(tgt)
        return tgt


# 3.4 完整Transformer模型（替换编码器为自定义CustomTransformerEncoder）
class Transformer(nn.Module):
    def __init__(self, src_vocab_size, tgt_vocab_size, seq_len, d_model=512,
                 dim_feedforward=2048, num_layers=6, nhead=8, dropout=0.1):
        super(Transformer, self).__init__()
        self.src_embedding = nn.Embedding(src_vocab_size, d_model)
        self.tgt_embedding = nn.Embedding(tgt_vocab_size, d_model)
        self.d_model = d_model
        self.pos_encoder = PositionalEncoding(d_model, max_len=seq_len, dropout=dropout)

        # -------------------------- 关键替换：用自定义CustomTransformerEncoder替代nn.TransformerEncoder --------------------------
        encoder_layer = EncoderLayer(d_model, nhead, dim_feedforward, dropout)
        self.encoder = CustomTransformerEncoder(encoder_layer, num_layers)  # 自定义编码器
        # --------------------------------------------------------------------------------------------------------------------

        # 解码器仍用PyTorch自带容器（若需自定义可同理实现CustomTransformerDecoder）
        decoder_layer = DecoderLayer(d_model, nhead, dim_feedforward, dropout)
        self.decoder = nn.TransformerDecoder(decoder_layer, num_layers=num_layers)

        self.fc_out = nn.Linear(d_model, tgt_vocab_size)
        self.init_weights()

    def init_weights(self):
        initrange = 0.1
        self.src_embedding.weight.data.uniform_(-initrange, initrange)
        self.tgt_embedding.weight.data.uniform_(-initrange, initrange)
        self.fc_out.bias.data.zero_()
        self.fc_out.weight.data.uniform_(-initrange, initrange)

    def generate_tgt_mask(self, tgt_seq_len):
        mask = nn.Transformer.generate_square_subsequent_mask(tgt_seq_len).to(Device)
        return mask  # (tgt_seq_len, tgt_seq_len)

    def forward(self, src, tgt):
        src_seq_len = src.size(1)
        tgt_seq_len = tgt.size(1)

        # 嵌入+位置编码（格式不变）
        src_emb = self.src_embedding(src).transpose(0, 1) * math.sqrt(self.d_model)  # (5, batch_size, 512)
        src_emb = self.pos_encoder(src_emb)
        tgt_emb = self.tgt_embedding(tgt).transpose(0, 1) * math.sqrt(self.d_model)  # (6, batch_size, 512)
        tgt_emb = self.pos_encoder(tgt_emb)

        # 解码器掩码（不变）
        tgt_mask = self.generate_tgt_mask(tgt_seq_len)

        # -------------------------- 自定义编码器前向传播（接口与原代码完全一致） --------------------------
        memory = self.encoder(src_emb, src_key_padding_mask=None)  # 输出格式仍为(5, batch_size, 512)
        # ----------------------------------------------------------------------------------------------------

        # 解码器逻辑不变（接收自定义编码器的memory）
        tgt_out = self.decoder(
            tgt_emb, memory,
            tgt_mask=tgt_mask,
            tgt_key_padding_mask=None,
            memory_key_padding_mask=None
        )

        output = self.fc_out(tgt_out).transpose(0, 1)  # (batch_size, 6, 11)
        return output


# -------------------------- 4. 训练函数（完全不变，兼容自定义编码器） --------------------------
def train():
    datas = create_data(batch_size=32, num=1000)

    # 模型初始化（参数不变，内部编码器已替换为自定义）
    model = Transformer(
        src_vocab_size=11,
        tgt_vocab_size=11,
        seq_len=10,
        d_model=512,
        dim_feedforward=4 * 512,
        num_layers=6,
        nhead=8
    ).to(Device)

    # 训练配置不变
    intervals = 30
    criterion = nn.CrossEntropyLoss()
    # opt = optim.SGD(model.parameters(), lr=1e-1)
    opt = optim.Adam(model.parameters(), lr=1e-4, betas=(0.9, 0.98), eps=1e-9)

    # 训练循环不变
    for interval in range(intervals):
        model.train()
        loss_all = 0.0
        for data, label in datas:
            data = data.to(Device)
            label = label.to(Device)

            input_label = label[:, :-1]  # (32,6)
            output_label = label[:, 1:]  # (32,6)

            opt.zero_grad()
            output = model(data, input_label)  # 调用自定义编码器的forward

            loss = criterion(
                output.reshape(-1, output.size(-1)),
                output_label.reshape(-1)
            )
            loss_all += loss.item()

            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            opt.step()

        avg_loss = loss_all / len(datas)
        print(f'第 {interval + 1} 轮, 平均损失: {avg_loss:.6f}, 学习率: {opt.param_groups[0]["lr"]:.8f}')

    # 测试逻辑不变（自定义编码器输出的memory可正常被解码器使用）
    with torch.no_grad():
        test_in_datas = [1, 2, 3, 4, 5]
        test_in_data = torch.tensor(test_in_datas, dtype=torch.long).view(1, -1)  # (1,5)
        label = torch.tensor([[0]], dtype=torch.long)  # 初始输入：起始符号0

        while True:
            test_in_data = test_in_data.to(Device)
            label = label.to(Device)

            # 生成解码器掩码（修复原代码掩码维度错误）
            tgt_seq_len = label.size(1)
            tgt_mask = model.generate_tgt_mask(tgt_seq_len)

            # 前向传播（自定义编码器参与计算）
            output = model(test_in_data, label)  # output: (1, tgt_seq_len, 11)
            log_softmax = F.log_softmax(output, dim=-1)
            max_token = log_softmax.argmax(dim=-1)  # (1, tgt_seq_len)

            # 拼接最新预测的token
            label = torch.cat([label, max_token[:, -1].unsqueeze(1)], dim=1)
            print(f"当前预测序列: {label.squeeze().tolist()}")

            # 遇到结束符9停止
            if max_token[:, -1].item() == 9:
                break


# -------------------------- 5. 启动训练 --------------------------
if __name__ == "__main__":
    train()